from .Utils import get_accessor_data

import pygltflib as gltf2
import numpy as np
# dirty patch for numpy high version
#np.bool = np.bool_

# convert numpy to vtk
from vtkmodules.util import (
    numpy_support
)

from vtkmodules.vtkCommonDataModel import (
    vtkPolyData,
)

from vtkmodules.vtkCommonCore import (
    vtkPoints,
    VTK_TYPE_UINT16,
    VTK_TYPE_INT32,
    VTK_TYPE_FLOAT32, 
)

from vtkmodules.vtkCommonDataModel import (
    vtkPolyData,
    vtkCellArray,
)

from vtkmodules.vtkFiltersHybrid import vtkWeightedTransformFilter 
from vtkmodules.vtkCommonTransforms import vtkTransform
from vtkmodules.vtkFiltersCore import vtkPolyDataTangents

from vtkmodules.vtkCommonExecutionModel import vtkPolyDataAlgorithm


# TODO: implement as vtk poly data algorithm...
class Primitive(vtkPolyDataAlgorithm):    
    def __init__(self) -> None:
        super().__init__()
        self.name = "Primitive"

        self.polyData = vtkPolyData()        
        self.tangents = None
        self.weightedTransformFilter = None

        self.material = None

    def Update(self):
        if self.tangents:
            self.tangents.Update()
        if self.weightedTransformFilter:
            self.weightedTransformFilter.Update()
        

    # parse Geometry from gltf model
    def ParseGeometry(self, gltf, prim):
        
        if prim.mode == gltf2.TRIANGLES:             
            # build poly data
            points  = vtkPoints()
            triangles = vtkCellArray()
            # index
            connectivity  = numpy_support.numpy_to_vtk(
                get_accessor_data(gltf, prim.indices),
                1, VTK_TYPE_INT32)
            offsets = numpy_support.numpy_to_vtk(
                np.arange(0, connectivity.GetSize()+1, 3, dtype=np.int32), # +1 for round error
                1, VTK_TYPE_INT32) 
            triangles.SetData(offsets, connectivity)

            # position
            position = numpy_support.numpy_to_vtk(
                get_accessor_data(gltf, prim.attributes.POSITION), 1, VTK_TYPE_FLOAT32)
            points.SetData(position)

            self.polyData.SetPoints(points)
            self.polyData.SetPolys(triangles)
            # normal
            if prim.attributes.NORMAL != None:
                normal = numpy_support.numpy_to_vtk(
                    get_accessor_data(gltf, prim.attributes.NORMAL), 1, VTK_TYPE_FLOAT32)
                self.polyData.GetPointData().SetNormals(normal)

            # only support texCoord0
            if prim.attributes.TEXCOORD_0 != None:
                texCoord = numpy_support.numpy_to_vtk(
                    get_accessor_data(gltf, prim.attributes.TEXCOORD_0), 1, VTK_TYPE_FLOAT32)
                self.polyData.GetPointData().SetTCoords(texCoord)
                
            # for skin animation
            if prim.attributes.JOINTS_0 != None and prim.attributes.WEIGHTS_0 != None:

                j = get_accessor_data(gltf, prim.attributes.JOINTS_0)
                j = j.astype(np.uint16)
                joints = numpy_support.numpy_to_vtk(j, 1, VTK_TYPE_UINT16)
                joints.SetName("JOINTS_0")

                w = get_accessor_data(gltf, prim.attributes.WEIGHTS_0)
                weights = numpy_support.numpy_to_vtk(w, 1, VTK_TYPE_FLOAT32)
                weights.SetName("WEIGHTS_0")
                
                self.polyData.GetPointData().AddArray(weights)
                self.polyData.GetPointData().AddArray(joints)

                # add weighted transform filter for skin animation
                self.weightedTransformFilter = vtkWeightedTransformFilter()
                self.weightedTransformFilter.SetInputData(self.polyData)
                self.weightedTransformFilter.SetWeightArray("WEIGHTS_0")
                #self.weightedTransformFilter.SetTransformIndexArray("JOINTS_0")
                                
                # initialize with default value...
                self.weightedTransformFilter.SetNumberOfTransforms(4)
                self.ResetSkinTransform()

                # generate tangents
            if prim.attributes.NORMAL != None and prim.attributes.TEXCOORD_0 != None:
                self.tangents = vtkPolyDataTangents()
                if self.weightedTransformFilter != None:                    
                    #self.weightedTransformFilter.Update()
                    self.tangents.SetInputConnection(self.weightedTransformFilter.GetOutputPort())
                else:
                    self.tangents.SetInputData(self.polyData)  
     
        else:
            print("Primitive(mode={}) is not implemented yet!".format(prim.mode)) 

        self.Update()

        if self.tangents:
            self.SetOutput(self.tangents.GetOutput())
        if self.weightedTransformFilter:
            self.SetOutput(self.weightedTransformFilter.GetPolyDataOutput())
        else:
            self.SetOutput(self.polyData)


    def UpdateSkinTransform(self, tm):        
        if(self.weightedTransformFilter):
            size = max(4, len(tm))
            self.weightedTransformFilter.SetNumberOfTransforms(size) 
            self.weightedTransformFilter.SetTransformIndexArray("JOINTS_0") 
            for i in range(len(tm)):
                self.weightedTransformFilter.SetTransform(tm[i], i)  
            self.Update()

    def ResetSkinTransform(self):
        if(self.weightedTransformFilter):
            size = self.weightedTransformFilter.GetNumberOfTransforms()
            identity = vtkTransform()
            for i in range(size):
                self.weightedTransformFilter.SetTransform(identity, i)  
            self.Update()
